/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package studio.raptor.ddal.core.executor.task;

import studio.raptor.ddal.core.connection.BackendConnection;
import studio.raptor.ddal.core.connection.ContextConnectionWrapper;
import studio.raptor.ddal.core.engine.StatementContext;
import studio.raptor.ddal.core.executor.BatchExecutionUnit;
import studio.raptor.ddal.core.executor.ExecutionUnit;
import studio.raptor.ddal.core.executor.resultset.ResultData;

import java.sql.SQLException;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * IUD语句执行任务
 *
 * @author Charley
 * @since 1.0
 */
public class IUDExecutionTask extends ExecutionTask {

    public IUDExecutionTask(Map<String, ContextConnectionWrapper> connectionsWrapper, ExecutionUnit executionUnit) {
        super(connectionsWrapper, executionUnit);
    }

    /**
     * 执行方法
     *
     * @return
     */
    @Override
    protected ResultData doExecute() throws SQLException {
        ResultData resultData=null;
        if (isPrepared) {
            resultData = this.executionUnit.getConnection().executePreparedUpdate(sql, parameters);
        } else {
            resultData = this.executionUnit.getConnection().executeUpdate(sql);
        }
        return resultData;
    }

    /**
     * 执行方法Batch
     */
    @Override
    protected ResultData doExecuteBatch(StatementContext statementContext) throws SQLException {
        ResultData resultData=null;
        String shardStr=this.executionUnit.getShard().toString();
        if (isPrepared) {
            if(isAddBatch){
                BackendConnection backendConnection=(BackendConnection)statementContext.getBackendConnectionMap().get(shardStr);
                BatchExecutionUnit batchExecutionUnit=statementContext.getBatchExecutionUnitMap().get(shardStr+this.executionUnit.getFinalSql().trim());
                resultData = backendConnection.executePreparedAddBatch(batchExecutionUnit,sql, parameters );
            }
            else if (isExecuteBatch){//遍历根据分区信息缓存的物理链接、辅助批处理执行单元
                resultData=executePreparedBatchToDB(statementContext);
            }
            else{
                resultData = this.executionUnit.getConnection().executePreparedUpdate(sql, parameters);
            }
        } else {
            if(isAddBatch){
                BackendConnection backendConnection=(BackendConnection)statementContext.getBackendConnectionMap().get(shardStr);
                BatchExecutionUnit batchExecutionUnit=statementContext.getBatchExecutionUnitMap().get(shardStr);
                resultData = backendConnection.executeAddBatch(batchExecutionUnit,sql );
            }
            else if (isExecuteBatch){
                resultData=executeBatchToDB(statementContext);
            }
            else{
                resultData = this.executionUnit.getConnection().executeUpdate(sql);
            }
        }
        statementContext.setMergedResult( resultData );
        return resultData;
    }

    private ResultData executeBatchToDB(StatementContext statementContext) throws SQLException{
        ResultData resultData=null;
        Map backendConnectionMap=statementContext.getBackendConnectionMap();
        Set<String> shardSet = backendConnectionMap.keySet();
        for (String shard:shardSet) {
            BackendConnection backendConnection=(BackendConnection)backendConnectionMap.get( shard );
            BatchExecutionUnit batchExecutionUnit=statementContext.getBatchExecutionUnitMap().get(shard );
            resultData=backendConnection.executeBatch(batchExecutionUnit);
        }
        //结果集合并，与入口处批处理sql提交顺序基本对不上
        BatchExecutionUnit batchExecutionUnitAll=new BatchExecutionUnit();
        for (String shard:shardSet){
            BatchExecutionUnit batchExecutionUnit=statementContext.getBatchExecutionUnitMap().get(shard );
            batchExecutionUnitAll.addAffectedRowsOfShard( batchExecutionUnit.calcBatchSqlAffectedRows() );
        }
        resultData=new ResultData(batchExecutionUnitAll.calcBatchSqlAffectedRows());
        return resultData;
    }

    private ResultData executePreparedBatchToDB(StatementContext statementContext) throws SQLException{
        ResultData resultData=null;
        Map backendConnectionMap=statementContext.getBackendConnectionMap();
        Set<String> shardSet = backendConnectionMap.keySet();
        for (String shard:shardSet) {
            BackendConnection backendConnection = (BackendConnection) backendConnectionMap.get( shard );
            List<String> finalSqlList=statementContext.getBatchSqlList();
            BatchExecutionUnit batchExecutionUnit;
            for (String finalSql:finalSqlList) {
                batchExecutionUnit = statementContext.getBatchExecutionUnitMap().get( shard+finalSql.trim());
                //非分表，其他分库上batchExecutionUnit会取不到（sql带用户名)
                if(null!=batchExecutionUnit){
                    //当前批处理sql(分库分表)
                    this.sql=finalSql;
                    resultData = backendConnection.executePreparedBatch( batchExecutionUnit );
                }
            }
        }
        //结果集合并，与入口处批处理sql提交顺序基本对不上
        BatchExecutionUnit batchExecutionUnitAll=new BatchExecutionUnit();
        for (String shard:shardSet){
            List<String> finalSqlList=statementContext.getBatchSqlList();
            for (String finalSql:finalSqlList) {
                BatchExecutionUnit batchExecutionUnit=statementContext.getBatchExecutionUnitMap().get(shard + finalSql.trim());
                if(null!=batchExecutionUnit){
                    batchExecutionUnitAll.addAffectedRowsOfShard( batchExecutionUnit.calcBatchSqlAffectedRows() );
                }
            }
        }
        resultData=new ResultData(batchExecutionUnitAll.calcBatchSqlAffectedRows());
        return resultData;
    }

}
